from contextlib import contextmanager
from selenium import webdriver
from selenium.webdriver.common.by import By
import platform

options = webdriver.ChromeOptions()
options.add_argument("--start-maximized")
options.add_argument('lang=zh_CN.UTF-8')
options.add_argument('--disable-extensions')
options.add_argument('--profile-directory=Default')
options.add_argument("--incognito")
options.add_argument("--disable-plugins-discovery")

if platform.system() == 'Windows':
    PATH = 'D:\\chrome\\chromedriver.exe'
else:
    # PATH = '/usr/bin/chromedriver'
    PATH = '/Users/luohao/app/chromedriver'


@contextmanager
def new_driver(timeout=60, headless=False, images=False, agent=None, proxy=None):
    if not images:
        prefs = {"profile.managed_default_content_settings.images": 2}
        options.add_experimental_option('prefs', prefs)

    if agent is not None:
        options.add_argument('user-agent="{}"'.format(agent))
    if proxy is not None:
        options.add_argument('--proxy-server={}'.format(proxy))

    if headless:
        options.add_argument('--headless')

    driver = webdriver.Chrome(executable_path=PATH, chrome_options=options)
    driver.delete_all_cookies()
    driver.set_page_load_timeout(timeout)

    try:
        yield driver
    finally:
        driver.close()
        driver.quit()


driverpool = []


@contextmanager
def keep_driver(timeout=600, headless=False, agent=None, proxy=None, images=False):
    if driverpool.__len__() > 0:
        driver = driverpool.pop()
    else:
        if not images:
            prefs = {"profile.managed_default_content_settings.images": 2}
            options.add_experimental_option('prefs', prefs)

        if agent is not None:
            options.add_argument('user-agent="{}"'.format(agent))
        if proxy is not None:
            options.add_argument('--proxy-server={}'.format(proxy))

        if headless:
            options.add_argument('--headless')
        driver = webdriver.Chrome(executable_path=PATH, chrome_options=options)
        driver.set_page_load_timeout(timeout)
    try:
        yield driver
    finally:
        driverpool.append(driver)
